import torch

def compute_plucker_rays(K, E, H, W):
    """
    K: [B, V, 3, 3]
    E: [B, V, 4, 4]
    """
    assert len(K.shape) == 4 and len(E.shape) == 4

    B, V, _, _ = K.shape
    device = K.device

    # make grid
    y_range = torch.arange(H, dtype = torch.float32, device = device)
    x_range = torch.arange(W, dtype = torch.float32, device = device)
    y_grid, x_grid = torch.meshgrid(y_range, x_range, indexing='ij')
    grid = torch.stack((x_grid, y_grid, torch.ones_like(x_grid)), dim = 2)
    grid = grid.unsqueeze(0).expand(B * V, -1, -1, -1)  # [B * V, H, W, 3]
    grid_reshaped = grid.permute(0, 3, 1, 2).reshape(B * V, 3, -1)


    # pixel to cam    
    K = K.reshape(B * V, 3, 3)
    E = E.reshape(B * V, 4, 4)

    K_inv = K.inverse()

    camera_coords = K_inv @ grid_reshaped

    # cam to world
    R = E[..., :3, :3]  # [B * V, 3, 3]
    T = E[..., :3, 3:]  # [B * V, 3, 1]

    dir_ray = R.mT @ camera_coords
    origin_world = -R.mT @ T

    # split and norm
    dir_ray = dir_ray.reshape(B, V, 3, -1)
    origin_world = origin_world.reshape(B, V, 3, -1)

    dir_ray = dir_ray / torch.norm(dir_ray, dim = -2, keepdim = True)

    m = torch.cross(origin_world, dir_ray, dim = -2)    
    out = torch.cat((dir_ray, m), dim = -2)
    out = out.reshape(B, V, 6, H, W)

    return out


def norm_cam(E, I):
    assert len(E.shape) == 4 and len(I.shape) == 4
    
    rots = E[:, :, :3, :3]   # [B, V, 3, 3]
    trans = E[:, :, :3, 3]   # [B, V, 3]

    # mean rots and trans
    mean_rot = rots.mean(dim = 1)  # [B, 3, 3]
    mean_trans = trans.mean(dim = 1, keepdim = True)  # [B, 3]

    # SVD
    U, S, Vt = torch.svd(mean_rot)
    mean_rot_normalized = U @ Vt.mT

    # mean trans
    trans_centered = trans - mean_trans  # [B, V, 3]

    # 归一化平移向量的长度
    trans_scale = torch.norm(trans_centered, dim = -1, keepdim = True)  # [B, 1, 3]
    scale, _ = trans_scale.max(dim = -2, keepdim = True) # [B, 1, 1]
    scale = scale + 1e-6

    trans_normalized = trans_centered / scale

    # 对于旋转，我们需要将每个旋转矩阵与平均旋转矩阵的逆相乘
    mean_rot_normalized = mean_rot_normalized.unsqueeze(1)
    rots_normalized = rots @ mean_rot_normalized.inverse()

    # 创建新的归一化外参矩阵
    E_normed = E.clone()
    I_normed = I.clone()

    E_normed[:, :, :3, :3] = rots_normalized
    E_normed[:, :, :3, 3] = trans_normalized

    I_normed[:, :, 0, 0] = I_normed[:, :, 0, 0] * scale.squeeze(-1)
    I_normed[:, :, 1, 1] = I_normed[:, :, 1, 1] * scale.squeeze(-1)

    return E_normed, I_normed